
from nnunet_dataloader.huanhu_ds import get_huanhu_loader
import numpy as np
import torch 
import torch.nn as nn 
from monai.networks.nets.basic_unet import BasicUNet
from light_training.evaluation.metric import dice
from light_training.utils.lr_scheduler import LinearWarmupCosineAnnealingLR
from medpy import metric
import glob 
import random 
torch.backends.cudnn.deterministic = True
random.seed(1)
torch.manual_seed(1)
torch.cuda.manual_seed(1)
np.random.seed(1)
import os
from tqdm import tqdm

os.environ["CUDA_VISIBLE_DEVICES"] = "1"

device = "cuda:0"
save_dir = "./logs_huanhu_nnunet/"

os.makedirs(save_dir, exist_ok=True)

data_paths = glob.glob("/home/xingzhaohu/sharefs/datasets/huanhu/resize_data/*.h5")
print("数据共：{} 例".format(len(data_paths)))

train_two_fold = []
test_two_fold = []
train_paths = []
test_paths = []

for i in range(78):
    train_paths.append(data_paths[i])
#
for i in range(78, len(data_paths)):
    test_paths.append(data_paths[i])

train_two_fold.append(train_paths)
test_two_fold.append(test_paths)

# ## 第二折
train_paths = []
test_paths = []
for i in range(78, len(data_paths)):
    train_paths.append(data_paths[i])

for i in range(78):
    test_paths.append(data_paths[i])

train_two_fold.append(train_paths)
test_two_fold.append(test_paths)

if __name__ == "__main__":
    fold = 1

    save_name = f"nnunet_model_e300_{fold}.pt"
    save_best_name = f"best_nnunet_model_e300_{fold}.pt"
    save_path = os.path.join(save_dir, save_name)
    save_best_path = os.path.join(save_dir, save_best_name)


    train_paths = train_two_fold[fold]
    val_paths = test_two_fold[fold]

    train_loader, val_loader = get_huanhu_loader(train_paths, val_paths)
    loss_func = nn.CrossEntropyLoss()
    model = BasicUNet(spatial_dims=3, 
                    in_channels=1, 
                    out_channels=2,
                    features=[16, 16, 32, 64, 128, 16])

    model.to(device)
    optimizer = torch.optim.Adam(model.parameters(), lr=1e-4, weight_decay=1e-3)
    scheduler = LinearWarmupCosineAnnealingLR(optimizer,
                                            warmup_epochs=10,
                                            max_epochs=200)
    report_loss = 0.0
    index = 0
    best_dice = 0

    for e in range(200):
        for i in tqdm(range(100), total=100):
            index += 1
            model.train()
            data = next(train_loader)
            image = data["data"].to(device)
            label = data["seg"].to(device)
            optimizer.zero_grad()

            image = nn.functional.interpolate(image, size=(32, 256, 256), mode="trilinear", align_corners=False)
            label = torch.unsqueeze(label, dim=1)
            label = nn.functional.interpolate(label, size=(32, 256, 256), mode="nearest")
            label = torch.squeeze(label, dim=1).long()

            pred = model(image)
            loss = loss_func(pred, label)
            report_loss += loss.item()

            loss.backward()
            optimizer.step()
            scheduler.step()

            if index % 500 == 0:
                print(f"epoch is {e}, loss is {report_loss}")
                report_loss = 0.0
                # validation 
                dices = []
                model.eval()
                for image, label in val_loader:
                    image = image.to(device)
                    label = label.to(device)

                    image = nn.functional.interpolate(image, size=(32, 256, 256), mode="trilinear", align_corners=False)
                    label = torch.unsqueeze(label, dim=1)
                    label = nn.functional.interpolate(label, size=(32, 256, 256), mode="nearest")
                    label = torch.squeeze(label, dim=1).long()

                    pred = model(image).argmax(dim=1).cpu().numpy()
                    label = label.cpu().numpy()
                    dices.append(metric.dc(pred, label))

                mean_dice = sum(dices) / len(dices)

                if best_dice < mean_dice:
                    best_dice = mean_dice

                    torch.save(model.state_dict(), save_best_path)
                
                print(f"mean dice is {mean_dice}, best dice is {best_dice}")


        torch.save(model.state_dict(), save_path)





    

